﻿namespace Microsoft.VisualStudio.TestTools.UnitTesting
{
    using System;
    using System.CodeDom.Compiler;
    using System.Collections.Generic;
    using System.Diagnostics.Contracts;
    using System.IO;
    using System.Reflection;
    using System.Resources;
    using Microsoft.CSharp;

    /// <summary>
    /// A utility class that helps dynamically generate assemblies.
    /// </summary>
    /// <remarks>
    /// This can be used to build assemblies for unit tests that need to build an assembly at run time.
    /// </remarks>
    public sealed class CompileHelper : IDisposable
    {
        private string assemblyFileName;
        private IList<string> assemblyReferences = new List<string>();
        private IList<string> classMembers = new List<string>();
        private string className;
        private CodeDomProvider codeProvider;
        private Assembly debugAssembly;
        private string debugAssemblyFileName;
        private Dictionary<string, Stream> embeddedResources = new Dictionary<string, Stream>();
        private string outputPath;
        private Assembly releaseAssembly;
        private string releaseAssemblyFileName;
        private List<string> sources = new List<string>();

        #region Constructors
        /// <summary>
        /// Initializes a new instance of the <see cref="CompileHelper"/> class.
        /// </summary>
        public CompileHelper()
        {
        }
        #endregion Constructors

        #region Public Properties
        /// <summary>
        /// Gets the name of the assembly file.
        /// </summary>
        /// <value>The name of the assembly file.</value>
        public string AssemblyFileName
        {
            get { return this.assemblyFileName; }
        }

        /// <summary>
        /// Gets the assembly references.
        /// </summary>
        /// <value>The assembly references.</value>
        public IList<string> AssemblyReferences
        {
            get { return this.assemblyReferences; }
        }

        /// <summary>
        /// Gets the generated debug assembly.
        /// </summary>
        /// <value>The generated debug assembly.</value>
        public Assembly DebugAssembly
        {
            get
            {
                if (this.debugAssembly == null)
                {
                    this.debugAssembly = Assembly.LoadFrom(this.DebugAssemblyFileName);
                }

                return this.debugAssembly;
            }
        }

        /// <summary>
        /// Gets the name of the debug assembly file.
        /// </summary>
        /// <value>The name of the debug assembly file.</value>
        public string DebugAssemblyFileName
        {
            get
            {
                if (this.debugAssemblyFileName == null)
                {
                    this.debugAssemblyFileName = this.Generate(true);
                }

                return this.debugAssemblyFileName;
            }
        }

        /// <summary>
        /// Gets the generated release assembly.
        /// </summary>
        /// <value>The generated release assembly.</value>
        public Assembly ReleaseAssembly
        {
            get
            {
                if (this.releaseAssembly == null)
                {
                    this.releaseAssembly = Assembly.LoadFrom(this.ReleaseAssemblyFileName);
                }

                return this.releaseAssembly;
            }
        }

        /// <summary>
        /// Gets the name of the release assembly file.
        /// </summary>
        /// <value>The name of the release assembly file.</value>
        public string ReleaseAssemblyFileName
        {
            get
            {
                if (this.releaseAssemblyFileName == null)
                {
                    this.releaseAssemblyFileName = this.Generate(false);
                }

                return this.releaseAssemblyFileName;
            }
        }
        #endregion Public Properties

        #region Public Methods
        /// <summary>
        /// Adds an assembly reference.
        /// </summary>
        /// <param name="assemblyReference">The assembly reference.</param>
        /// <returns>This instance.</returns>
        public CompileHelper AddAssemblyReference(string assemblyReference)
        {
            Contract.Requires(!string.IsNullOrWhiteSpace(assemblyReference));

            this.assemblyReferences.Add(assemblyReference);

            return this;
        }

        /// <summary>
        /// Adds the assembly reference.
        /// </summary>
        /// <param name="assembly">The assembly to reference.</param>
        /// <returns>This instance.</returns>
        public CompileHelper AddAssemblyReference(Assembly assembly)
        {
            if (assembly == null)
            {
                throw new ArgumentNullException("assembly");
            }

            Contract.EndContractBlock();

            this.assemblyReferences.Add(assembly.Location);

            return this;
        }

        /// <summary>
        /// Starts adding a new class.
        /// </summary>
        /// <param name="className">Name of the class.</param>
        /// <returns>This instance.</returns>
        /// <remarks>
        /// Call this method and then call <see cref="AddMember"/> to add members to the class.
        /// </remarks>
        public CompileHelper AddClass(string className)
        {
            this.EnsureClassMembersAreAdded();

            this.className = className;

            return this;
        }

        /// <summary>
        /// Adds an embedded resource.
        /// </summary>
        /// <param name="resourceName">Name of the embedded resource. Note that ".resources" is appended.</param>
        /// <param name="writerAction">Builds the embedded resource.</param>
        /// <returns>This instance.</returns>
        /// <remarks>
        /// Each resource is generated by calling the <paramref name="writerAction"/>.
        /// </remarks>
        public CompileHelper AddEmbeddedResource(string resourceName, Action<ResourceWriter> writerAction)
        {
            Contract.Requires(resourceName != null);
            Contract.Requires(writerAction != null);

            var resourceStream = new MemoryStream();
            using (var writer = new ResourceWriter(resourceStream))
            {
                writerAction(writer);
                writer.Generate();

                // The original resourceStream will be closed when the ResourceWriter is disposed so copy to new stream.
                resourceStream = new MemoryStream(
                    resourceStream.GetBuffer(),
                    0,
                    (int)resourceStream.Length);
            }

            this.embeddedResources.Add(resourceName, resourceStream);

            return this;
        }

        /// <summary>
        /// Adds the member to the current class.
        /// </summary>
        /// <param name="member">The member to add.</param>
        /// <returns>This instance.</returns>
        public CompileHelper AddMember(string member)
        {
            Contract.Requires(!string.IsNullOrWhiteSpace(member));

            // Default to the class name Class if none was specified.
            if (this.className == null)
            {
                this.className = "Class";
            }

            this.classMembers.Add(member);

            return this;
        }

        /// <summary>
        /// Adds source code to compile.
        /// </summary>
        /// <param name="source">The source code.</param>
        /// <returns>This instance.</returns>
        public CompileHelper AddSource(string source)
        {
            Contract.Requires(!string.IsNullOrWhiteSpace(source));

            this.sources.Add(source);

            return this;
        }

        /// <summary>
        /// Adds an assembly reference that the user will provide the full path to.
        /// </summary>
        /// <param name="type">The type.</param>
        /// <param name="assemblyReference">The assembly reference.</param>
        /// <returns>This instance.</returns>
        public CompileHelper AddUserSuppliedAssemblyReference(Type type, string assemblyReference)
        {
            Contract.Requires(!string.IsNullOrWhiteSpace(assemblyReference));

            var assemblyFullPath = Microsoft.VisualStudio.TestTools.UnitTesting.AssemblyReferences.GetReferenceFullPath(type, assemblyReference);
            this.assemblyReferences.Add(assemblyFullPath);

            if (this.outputPath == null)
            {
                this.outputPath = Path.GetDirectoryName(assemblyFullPath);
            }

            return this;
        }

        /// <summary>
        /// Performs application-defined tasks associated with freeing, releasing, or resetting unmanaged resources.
        /// </summary>
        public void Dispose()
        {
            if (this.codeProvider != null)
            {
                this.codeProvider.Dispose();
                this.codeProvider = null;
            }
        }

        /// <summary>
        /// Generates the assembly.
        /// </summary>
        /// <param name="debugMode"><c>true</c> if the assembly should be compiled in debug mode.</param>
        /// <returns>The file name of the generated assembly.</returns>
        /// <exception cref="InvalidOperationException">A compiler error was found.</exception>
        [System.Diagnostics.CodeAnalysis.SuppressMessage("Microsoft.Reliability", "CA2000:Dispose objects before losing scope", Justification = "False positive.")]
        public string Generate(bool debugMode)
        {
            this.EnsureClassMembersAreAdded();

            // If a file name hasn't been specified then build a random name for the assembly.
            if (string.IsNullOrEmpty(this.assemblyFileName))
            {
                this.assemblyFileName = this.GetTemporaryFileName() + ".dll";
            }

            // Make sure we haven't already generated this assembly.
            if (File.Exists(this.assemblyFileName))
            {
                File.Delete(this.assemblyFileName);
            }

            // Compile the code into the dynamically built assembly.
            var options = new CompilerParameters()
            {
                GenerateExecutable = false,
                GenerateInMemory = false,
                IncludeDebugInformation = debugMode,
                OutputAssembly = this.assemblyFileName,
            };

            // Add in common assembly references.
            options.ReferencedAssemblies.Add("System.dll");

            // Add in any assembly references.
            foreach (string assemblyReference in this.assemblyReferences)
            {
                options.ReferencedAssemblies.Add(assemblyReference);
            }

            // Generate the embedded resource if any.
            foreach (var entry in this.embeddedResources)
            {
                var resourceFileName = this.GetFileNameAndEnsureDirectoryExists(entry.Key)
                    + ".resources";
                using (var fileStream = new FileStream(resourceFileName, FileMode.Create))
                {
                    entry.Value.CopyTo(fileStream);
                }

                options.EmbeddedResources.Add(resourceFileName);
            }

            // If a code provider wasn't provided then default to CSharp.
            if (this.codeProvider == null)
            {
                this.codeProvider = new CSharpCodeProvider();
            }

            var compilerResults = this.codeProvider.CompileAssemblyFromSource(options, this.sources.ToArray());

            // If there were any compiler errors throw an exception.
            if (compilerResults.Errors.Count > 0)
            {
                throw new InvalidOperationException("Compile Error: " + compilerResults.Errors[0].ErrorText);
            }

            return this.assemblyFileName;
        }

        /// <summary>
        /// Sets the name of the assembly file.
        /// </summary>
        /// <param name="assemblyFileName">Name of the assembly file.</param>
        /// <returns>This instance.</returns>
        [System.Diagnostics.CodeAnalysis.SuppressMessage("Microsoft.Maintainability", "CA1500:VariableNamesShouldNotMatchFieldNames", MessageId = "assemblyFileName", Justification = "By desing.")]
        public CompileHelper SetAssemblyFileName(string assemblyFileName)
        {
            this.assemblyFileName = assemblyFileName;

            return this;
        }

        /// <summary>
        /// Sets the code provider.
        /// </summary>
        /// <param name="codeProviderName">Name of the code provider.</param>
        /// <returns>This instance.</returns>
        public CompileHelper SetCodeProvider(string codeProviderName)
        {
            Contract.Requires((codeProviderName == "cs") || (codeProviderName == "vb"));

            if ("cs".Equals(codeProviderName, StringComparison.OrdinalIgnoreCase))
            {
                this.SetCodeProvider(new CSharpCodeProvider());
            }
            else if ("vb".Equals(codeProviderName, StringComparison.OrdinalIgnoreCase))
            {
                this.SetCodeProvider(new VisualBasic.VBCodeProvider());
            }
            else
            {
                throw new ArgumentException("The name must be 'cs' or 'vb'.", "codeProviderName");
            }

            return this;
        }

        /// <summary>
        /// Sets the code provider.
        /// </summary>
        /// <param name="codeProvider">The code provider.</param>
        /// <returns>This instance.</returns>
        [System.Diagnostics.CodeAnalysis.SuppressMessage("Microsoft.Maintainability", "CA1500:VariableNamesShouldNotMatchFieldNames", MessageId = "codeProvider", Justification = "By design.")]
        public CompileHelper SetCodeProvider(CodeDomProvider codeProvider)
        {
            Contract.Requires(codeProvider != null);

            this.codeProvider = codeProvider;

            return this;
        }

        /// <summary>
        /// Sets the output path.
        /// </summary>
        /// <param name="path">The path.</param>
        /// <returns>This instance.</returns>
        public CompileHelper SetOutputPath(string path)
        {
            this.outputPath = path;

            return this;
        }
        #endregion Public Methods

        #region Private Methods
        private string GetFileNameAndEnsureDirectoryExists(string name)
        {
            var directoryName = Path.GetDirectoryName(name);
            if (!string.IsNullOrEmpty(directoryName))
            {
                var fullDirectoryPath = Path.Combine(Environment.CurrentDirectory, directoryName);
                if (!Directory.Exists(fullDirectoryPath))
                {
                    Directory.CreateDirectory(fullDirectoryPath);
                }
            }

            return Path.Combine(Environment.CurrentDirectory, name);
        }

        private string GetTemporaryFileName()
        {
            var path = this.outputPath ?? Environment.CurrentDirectory;
            return Path.Combine(path, "A" + Guid.NewGuid().ToString("N"));
        }

        private void EnsureClassMembersAreAdded()
        {
            if (this.className != null)
            {
                this.AddSource("public class " + this.className + @"
{"
                    + string.Join(Environment.NewLine, this.classMembers) + @"
}");

                this.className = null;
            }

            this.classMembers.Clear();
        }
        #endregion Private Methods
    }
}